# -*- coding: gb2312 -*-
from pymongo import MongoClient
from WindPy import *

import time
import math

from StockDailyFactorGeneralFunctions import StockDailyFactorGeneralFunctions

class DateGeneralFunctions:
    def __init__(self, colSysData = 'sys_data'):


        funcObj = StockDailyFactorGeneralFunctions()
        self.fillType = funcObj.fillType
        self.DBIP = funcObj.DBIP
        self. DBPort = funcObj.DBPort
        self.DBName = funcObj.DBName


        self.stockIdList = []
        self.periodTransDateList = []
        self.colSysData = colSysData
        w.start()

    def retrieveAllTransDateList(self, startDateStr, endDateStr):
        resTDates = w.tdays(startDateStr, endDateStr)

        allTDatesList = []

        try:
            for tDatesStr in resTDates.Data[0]:
                allTDatesList.append(tDatesStr.strftime('%Y%m%d'))
        except IndexError:
            print('Transaction date is up to date already.')
            return allTDatesList

        return allTDatesList

    def retrievePeriodEndTransDateList(self, startDateStr, endDateStr, periodTypeStr):
        resTDates = w.tdays(startDateStr, endDateStr, 'Period=' + periodTypeStr)

        periodEndTDatesList = []

        for tDatesStr in resTDates.Data[0]:
            periodEndTDatesList.append(tDatesStr.strftime('%Y%m%d'))

        return periodEndTDatesList

#    def setDatesType(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d') ):
    def setDatesType(self, startDateStr, endDateStr):
        allTDateList = self.retrieveAllTransDateList(startDateStr, endDateStr)
        weekEndDateList = self.retrievePeriodEndTransDateList(startDateStr, endDateStr, 'W')
        monthEndDateList = self.retrievePeriodEndTransDateList(startDateStr, endDateStr, 'M')
        quarterEndDateList = self.retrievePeriodEndTransDateList(startDateStr, endDateStr, 'Q')

        aDateDict = {}

        for oneTDateStr in allTDateList:
            aDateDict[oneTDateStr] = ''
            if oneTDateStr in weekEndDateList:
                aDateDict[oneTDateStr] = aDateDict[oneTDateStr] + 'W'
            if oneTDateStr in monthEndDateList:
                aDateDict[oneTDateStr] = aDateDict[oneTDateStr] + 'M'
            if oneTDateStr in quarterEndDateList:
                aDateDict[oneTDateStr] = aDateDict[oneTDateStr] + 'Q'

        return aDateDict

    def writeBasePeriodEndTransDatesToMongoDB(self, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr, startDateStr, endDateStr):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collection_data = db[colNameStr]
        aData = []
        aData.append(self.setDatesType(startDateStr, endDateStr))
        collection_data.insert(aData)

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.updateDataUpdateDate(dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr, currentDateStr=endDateStr)


        print('Finished writing base ' + dbUpdateDateFieldAliasStr + 'to database.')

    def updatePeriodEndTransDatesToMongoDB(self, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr, endDateStr = time.strftime('%Y-%m-%d')):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        client = MongoClient(DBIP, DBPort)

        funcObj = StockDailyFactorGeneralFunctions()

        startDateStr = funcObj.retrieveDataUpdateDate(dbUpdateDateFiledStr)

        db = client[DBName]

        collection_data = db[colNameStr]

        aData = []
        aData.append(self.setDatesType(startDateStr, endDateStr))


        dateDict = aData[0]


        for (k, v) in dateDict.items():
           collection_data.update({}, {'$set': {k: v}}, upsert=True)

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.updateDataUpdateDate(dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr, currentDateStr=endDateStr)


        print('Finished updating' + dbUpdateDateFieldAliasStr + 'to database.')

#obj = DateGeneralFunctions()




#obj.writeBasePeriodEndTransDatesToMongoDB('trans_dates', 'trans_dates_update_date', 'transaction days', '2012-01-01', '2017-05-01')
#obj.updatePeriodEndTransDatesToMongoDB('trans_dates', 'trans_dates_update_date', 'transaction days')